import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.optim as optim
from sklearn.metrics import accuracy_score
from torchvision.models import resnet18
from experiments.experiments1 import get_is_correct_arr, count_elements_above_threshold, save_list_to_file
import tensorflow as tf

# 检查是否有可用的 GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 数据预处理
transform = transforms.Compose([
    transforms.Grayscale(num_output_channels=3),  # 将灰度图像转换为3通道图像
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# 下载和加载 Fashion-MNIST 数据集
trainset = torchvision.datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)

testset = torchvision.datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)

print(f"训练集数据量为{len(trainset)}条")
print(f"测试集数据量为{len(testset)}条")


# 定义 ResNet-18 模型
class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18, self).__init__()
        self.model = resnet18(pretrained=False, num_classes=10)  # Fashion-MNIST 有 10 个类别

    def forward(self, x):
        return self.model(x)


# 实例化模型、损失函数和优化器
net = ResNet18().to(device)  # 将模型移动到 GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

# 设置训练参数
epochs = 200
threshold = 30
threshold_value = 30

predicted_result = []
# 训练模型
for epoch in range(epochs):
    # 训练
    running_loss = 0.0
    net.train()
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到 GPU
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    # 测试模型
    correct = 0
    total = 0
    net.eval()

    pre_labels = torch.tensor([]).to(device)  # 将预测结果初始化在 GPU 上
    true_labels = torch.tensor([]).to(device)
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)  # 将数据移动到 GPU
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            # 保存预测的结果
            pre_labels = torch.cat((pre_labels, predicted))
            true_labels = torch.cat((true_labels, labels))

    is_correct_arr = get_is_correct_arr(true_labels, pre_labels)
    predicted_result.append(is_correct_arr)
    # 计算准确率
    accuracy = accuracy_score(true_labels.cpu(), pre_labels.cpu())  # 将张量移动到 CPU 以计算准确率
    print(f'Accuracy of the network on the 10000 test images: {accuracy:.2f}%')

predicted_result = tf.convert_to_tensor(predicted_result, dtype=tf.int32)
# 计算可以忽略的比例
ignore_rate_list = []
for i in range(1, epochs+1):
    if i < threshold:
        ignore_rate_list.append(0)
        continue
    span_tensor = predicted_result[i-30:i]
    sum_along_columns = tf.reduce_sum(span_tensor, axis=0)
    ignore_num = count_elements_above_threshold(sum_along_columns, threshold_value)
    ignore_rate = ignore_num / tf.size(sum_along_columns).numpy() * 100
    ignore_rate_list.append(ignore_rate)

print(ignore_rate_list)
save_list_to_file(ignore_rate_list, "./data_result/fashion_mnist.txt")
